/**
 * 
 */
package edu.umd.clip.lm.tools;

import java.io.*;

import com.meetup.memcached.MemcachedClient;
import com.meetup.memcached.SockIOPool;

import edu.berkeley.nlp.util.*;
import edu.umd.clip.jobs.JobManager;
import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.decoding.NgramMarginalization;
import edu.umd.clip.lm.util.IO;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class LoadDumpedNgrams {
	public static class Options {
        @Option(name = "-config", required = true, usage = "XML config file")
		public String config;
        @Option(name = "-input", required = false, usage = "Evaluation data file (Default: stdin)")
		public String input;
        @Option(name = "-jobs", usage = "number of concurrent jobs (default: 1)")
        public int jobs = 1;
        @Option(name = "-debug", usage = "decoder debug level (default: 0)")
        public int debug = 0;
        @Option(name = "-lc", usage = "lowercase the input")
        public boolean lower = false;
        @Option(name = "-uc", usage = "uppercase the input")
        public boolean upper = false;
        @Option(name = "-memcached", required = true, usage = "Memcached server address in the format 'hostname:port'")
		public String memcached;    
	}

	/**
	 * @param args
	 * @throws IOException 
	 */
	public static void main(String[] args) throws IOException {
        OptionParser optParser = new OptionParser(Options.class);
        final Options opts = (Options) optParser.parse(args, true);

        if (opts.upper && opts.lower) {
        	System.err.println("Lowercase AND uppercase? doesn't sound right.");
        	return;
        }
        
        JobManager.initialize(opts.jobs);
		Thread thread = new Thread(JobManager.getInstance(), "Job Manager");
		thread.setDaemon(true);
		thread.start();
       
		Experiment.initialize(opts.config);
        Experiment experiment = Experiment.getInstance();
		final InputParser parser = new PlainInputParser(experiment.getTupleDescription());
		
		final MemcachedClient memcached;
		{
			SockIOPool sockPool = SockIOPool.getInstance();
			sockPool.setServers(new String[] {opts.memcached});
			sockPool.setInitConn(JobManager.getInstance().getNumWorkers());
			sockPool.initialize();
			
			memcached = new MemcachedClient();
			memcached.setPrimitiveAsString(false);
			memcached.setSanitizeKeys(false);
			memcached.setCompressEnable(false);
		}
		
		BufferedReader reader = IO.getReader(opts.input == null ? System.in : IO.getInputStream((opts.input)));

	    for(String line = reader.readLine(); line != null; line = reader.readLine()) {
	    	Pair<long[], Double> pair = NgramMarginalization.readDumpedNgram(parser, line);
	    	long ngram[] = pair.getFirst();
	    	Double prob = pair.getSecond();
	    	String ngramKey = NgramMarginalization.ngram2memcachedKey(ngram);
	    	memcached.set(ngramKey, prob);
	    }
	}
}
